from __future__ import annotations

import ast
import os
from dataclasses import dataclass, field
from pathlib import Path
from textwrap import dedent
from typing import Dict, List, Optional

import libcst as cst
# imports
import pytest
from codeflash.code_utils.edit_generated_tests import \
    add_runtime_comments_to_generated_tests

# --- Minimal stubs for dependencies ---

# InvocationId stub
@dataclass(frozen=True)
class InvocationId:
    test_module_path: str
    test_class_name: str
    test_function_name: str
    iteration_id: str

# GeneratedTests stub
@dataclass
class GeneratedTests:
    generated_original_test_source: str
    instrumented_behavior_test_source: str
    instrumented_perf_test_source: str
    behavior_file_path: Path
    perf_file_path: Path

# GeneratedTestsList stub
@dataclass
class GeneratedTestsList:
    generated_tests: List[GeneratedTests]

# TestConfig stub
@dataclass
class TestConfig:
    tests_root: Path
    project_root_path: Path
from codeflash.code_utils.edit_generated_tests import \
    add_runtime_comments_to_generated_tests

# ========== UNIT TESTS ==========

# Helper to build a minimal test source with a function call
def make_test_source(func_name: str, call_line: str = None, class_name: Optional[str] = None) -> str:
    if class_name:
        return f"""
class {class_name}:
    def test_func(self):
        codeflash_output = {func_name}()
"""
    else:
        return f"""
def test_func():
    codeflash_output = {func_name}()
"""

# Helper to create a GeneratedTestsList with one test
def make_generated_tests_list(source: str, behavior_path: Path, perf_path: Path) -> GeneratedTestsList:
    return GeneratedTestsList([
        GeneratedTests(
            generated_original_test_source=source,
            instrumented_behavior_test_source="",
            instrumented_perf_test_source="",
            behavior_file_path=behavior_path,
            perf_file_path=perf_path,
        )
    ])

# Helper to create InvocationId for a function
def make_invocation_id(test_module_path, test_class_name, test_function_name, iteration_id):
    return InvocationId(
        test_module_path=test_module_path,
        test_class_name=test_class_name,
        test_function_name=test_function_name,
        iteration_id=iteration_id
    )

# ========== BASIC TEST CASES ==========

def test_basic_runtime_comment_inserted():
    """Basic: Should add comment for a single function call in a function."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    # Matching InvocationId
    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    # Simulate original and optimized runtimes
    original_runtimes = {inv_id: [200_000_000]}  # 200ms
    optimized_runtimes = {inv_id: [100_000_000]} # 100ms

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    # The output should contain the runtime comment after the function call
    out_source = result.generated_tests[0].generated_original_test_source

def test_basic_runtime_comment_slower():
    """Basic: Should indicate 'slower' if optimized is slower."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}  # 100ms
    optimized_runtimes = {inv_id: [200_000_000]} # 200ms

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_basic_no_matching_invocation():
    """Basic: If no matching invocation, no comment should be added."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    # InvocationId does not match function name
    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="other_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}
    optimized_runtimes = {inv_id: [80_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_basic_multiple_calls():
    """Basic: Multiple calls in one function should each get a comment."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = """
def test_func():
    codeflash_output = myfunc()
    codeflash_output = myfunc()
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    # Two invocation ids, one for each call (simulate iteration_id as 0 and 1)
    inv_id0 = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    inv_id1 = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="1_0"
    )
    original_runtimes = {inv_id0: [100_000_000], inv_id1: [200_000_000]}
    optimized_runtimes = {inv_id0: [80_000_000], inv_id1: [100_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_basic_class_method():
    """Basic: Should handle test method in a class."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name, class_name="TestClass")
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="TestClass",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}
    optimized_runtimes = {inv_id: [90_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

# ========== EDGE TEST CASES ==========

def test_edge_zero_runtimes():
    """Edge: If either runtime is zero, no comment should be added."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    # One runtime is zero
    original_runtimes = {inv_id: [0]}
    optimized_runtimes = {inv_id: [100_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_edge_no_codeflash_output():
    """Edge: If there is no codeflash_output assignment, no comment should be added."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = """
def test_func():
    pass
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    original_runtimes = {}
    optimized_runtimes = {}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_edge_multiple_classes():
    """Edge: Handles multiple classes and functions with same name in different classes."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = """
class TestClassA:
    def test_func(self):
        codeflash_output = myfunc()
class TestClassB:
    def test_func(self):
        codeflash_output = myfunc()
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id_a = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="TestClassA",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    inv_id_b = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="TestClassB",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id_a: [100_000_000], inv_id_b: [200_000_000]}
    optimized_runtimes = {inv_id_a: [80_000_000], inv_id_b: [100_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_edge_non_matching_path():
    """Edge: If the test file path does not match InvocationId, no comment should be added."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = make_test_source(func_name)
    behavior_path = Path("tests/other_module.py")
    perf_path = Path("tests/other_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}
    optimized_runtimes = {inv_id: [80_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source

def test_edge_attribute_call():
    """Edge: Should detect calls like obj.myfunc() as well."""
    func_name = "myfunc"
    qualified_name = func_name
    test_source = """
def test_func():
    codeflash_output = obj.myfunc()
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    inv_id = make_invocation_id(
        test_module_path="tests.test_module",
        test_class_name="",
        test_function_name="test_func",
        iteration_id="0_0"
    )
    original_runtimes = {inv_id: [100_000_000]}
    optimized_runtimes = {inv_id: [80_000_000]}

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source


def test_large_many_tests(monkeypatch):
    """Large Scale: Handles many test files and calls efficiently."""
    func_name = "myfunc"
    qualified_name = func_name
    N = 100  # 100 test functions
    test_sources = []
    behavior_paths = []
    perf_paths = []
    original_runtimes = {}
    optimized_runtimes = {}
    tests = []

    for i in range(N):
        source = f"""
def test_func_{i}():
    codeflash_output = myfunc()
"""
        behavior_path = Path(f"tests/test_module_{i}.py")
        perf_path = Path(f"tests/test_module_{i}.py")
        tests.append(GeneratedTests(
            generated_original_test_source=source,
            instrumented_behavior_test_source="",
            instrumented_perf_test_source="",
            behavior_file_path=behavior_path,
            perf_file_path=perf_path,
        ))
        inv_id = make_invocation_id(
            test_module_path=f"tests.test_module_{i}",
            test_class_name="",
            test_function_name=f"test_func_{i}",
            iteration_id="0_0"
        )
        # Each test has unique timings
        original_runtimes[inv_id] = [100_000_000 + i*1_000_000]
        optimized_runtimes[inv_id] = [90_000_000 + i*1_000_000]

    tests_list = GeneratedTestsList(tests)
    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    # Each output should have correct comment
    for i in range(N):
        out_source = result.generated_tests[i].generated_original_test_source
        expected_comment = f"# {100+i}ms -> {90+i}ms (10.00% faster)"

def test_large_many_calls_per_test():
    """Large Scale: Handles many calls in a single test function."""
    func_name = "myfunc"
    qualified_name = func_name
    N = 100  # 100 calls in one function
    call_lines = "\n".join([f"    codeflash_output = myfunc()" for _ in range(N)])
    test_source = f"""
def test_func():
{call_lines}
"""
    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)

    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(N):
        inv_id = make_invocation_id(
            test_module_path="tests.test_module",
            test_class_name="",
            test_function_name="test_func",
            iteration_id=f"{i}_0"
        )
        original_runtimes[inv_id] = [100_000_000 + i*1_000_000]
        optimized_runtimes[inv_id] = [90_000_000 + i*1_000_000]

    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source
    # All 100 comments should appear
    for i in range(N):
        expected_comment = f"# {100+i}ms -> {90+i}ms (10.00% faster)"

def test_large_many_classes_and_methods():
    """Large Scale: Handles many classes and methods."""
    func_name = "myfunc"
    qualified_name = func_name
    N = 50  # 50 classes, each with 2 methods
    test_source = ""
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(N):
        test_source += f"""
class TestClass{i}:
    def test_func_a(self):
        codeflash_output = myfunc()
    def test_func_b(self):
        codeflash_output = myfunc()
"""
        inv_id_a = make_invocation_id(
            test_module_path=f"tests.test_module",
            test_class_name=f"TestClass{i}",
            test_function_name="test_func_a",
            iteration_id="0_0"
        )
        inv_id_b = make_invocation_id(
            test_module_path=f"tests.test_module",
            test_class_name=f"TestClass{i}",
            test_function_name="test_func_b",
            iteration_id="0_0"
        )
        original_runtimes[inv_id_a] = [100_000_000 + i*1_000_000]
        optimized_runtimes[inv_id_a] = [90_000_000 + i*1_000_000]
        original_runtimes[inv_id_b] = [200_000_000 + i*1_000_000]
        optimized_runtimes[inv_id_b] = [180_000_000 + i*1_000_000]

    behavior_path = Path("tests/test_module.py")
    perf_path = Path("tests/test_module.py")
    tests_list = make_generated_tests_list(test_source, behavior_path, perf_path)
    test_cfg = TestConfig(tests_root=Path("tests"), project_root_path=Path("."))

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name, test_cfg, tests_list, original_runtimes, optimized_runtimes
    ); result = codeflash_output
    out_source = result.generated_tests[0].generated_original_test_source
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.

import ast
import os
from dataclasses import dataclass, field
from pathlib import Path
from textwrap import dedent
from typing import Dict, List

# Use libcst for parsing and transforming code
import libcst as cst
# imports
import pytest
from codeflash.code_utils.edit_generated_tests import \
    add_runtime_comments_to_generated_tests

# --- Minimal stub implementations for dependencies ---

# Simulate InvocationId
@dataclass(frozen=True)
class InvocationId:
    test_module_path: str
    test_class_name: str
    test_function_name: str
    iteration_id: str

# Simulate GeneratedTests
@dataclass
class GeneratedTests:
    generated_original_test_source: str
    instrumented_behavior_test_source: str = ""
    instrumented_perf_test_source: str = ""
    behavior_file_path: Path = Path("tests/test_file.py")
    perf_file_path: Path = Path("tests/test_file.py")

# Simulate GeneratedTestsList
@dataclass
class GeneratedTestsList:
    generated_tests: List[GeneratedTests]

# Simulate TestConfig
@dataclass
class TestConfig:
    tests_root: Path
    project_root_path: Path
from codeflash.code_utils.edit_generated_tests import \
    add_runtime_comments_to_generated_tests


def make_file_paths():
    project_root = Path("/project")
    tests_root = project_root / "tests"
    test_file = tests_root / "test_file.py"
    return project_root, tests_root, test_file

# Helper to build InvocationId
def make_invocation_id(
    test_module_path="tests.test_file",
    test_class_name="TestClass",
    test_function_name="test_func",
    iteration_id="0_0"
):
    return InvocationId(
        test_module_path=test_module_path,
        test_class_name=test_class_name,
        test_function_name=test_function_name,
        iteration_id=iteration_id
    )

# --- BASIC TEST CASES ---

def test_single_function_single_call_adds_comment():
    """
    Basic: Single function, single call, matching runtime data.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(1)
        codeflash_output = result
"""
    # Setup test config and test object
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])

    # Setup runtime data
    invocation_id = make_invocation_id()
    original_runtimes = {invocation_id: [1000]}
    optimized_runtimes = {invocation_id: [500]}

    # Call function
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output

    # Assert: comment is present and correct
    out_code = result.generated_tests[0].generated_original_test_source

def test_no_matching_invocation_no_comment():
    """
    Basic: No matching invocation id, no comment is added.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(2)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])

    # Empty runtime data
    original_runtimes = {}
    optimized_runtimes = {}

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_multiple_calls_multiple_comments():
    """
    Basic: Multiple calls in one function, multiple comments.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result1 = myfunc(1)
        codeflash_output = result1
        result2 = myfunc(2)
        codeflash_output = result2
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])

    # Two calls, two invocation ids
    invocation_id1 = make_invocation_id(iteration_id="0_0")
    invocation_id2 = make_invocation_id(iteration_id="1_0")
    original_runtimes = {invocation_id1: [1000], invocation_id2: [2000]}
    optimized_runtimes = {invocation_id1: [800], invocation_id2: [1500]}

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_comment_shows_slower_when_optimized_slower():
    """
    Basic: Optimized is slower, comment says 'slower'.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(5)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])

    invocation_id = make_invocation_id()
    original_runtimes = {invocation_id: [1000]}
    optimized_runtimes = {invocation_id: [1500]}

    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

# --- EDGE TEST CASES ---

def test_function_with_no_calls_no_comment():
    """
    Edge: Function does not call the target function at all.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        x = 1 + 2
        y = x * 2
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes={},
        optimized_runtimes={}
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_zero_runtimes_no_comment():
    """
    Edge: Runtimes are zero, comment is not added.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(1)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    invocation_id = make_invocation_id()
    original_runtimes = {invocation_id: [0]}
    optimized_runtimes = {invocation_id: [0]}
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_function_with_attribute_call():
    """
    Edge: Function call is as an attribute (obj.myfunc), comment is added.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        obj = SomeClass()
        result = obj.myfunc(7)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    invocation_id = make_invocation_id()
    original_runtimes = {invocation_id: [100]}
    optimized_runtimes = {invocation_id: [50]}
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_multiple_classes_and_functions():
    """
    Edge: Multiple classes and functions, only correct function gets comment.
    """
    project_root, tests_root, test_file = make_file_paths()
    test_code = """
class TestClass:
    def test_func(self):
        result = myfunc(1)
        codeflash_output = result

class OtherClass:
    def test_func(self):
        result = myfunc(2)
        codeflash_output = result
"""
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    # Only first class/function gets runtime data
    invocation_id = make_invocation_id(test_class_name="TestClass")
    original_runtimes = {invocation_id: [1000]}
    optimized_runtimes = {invocation_id: [900]}
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source


def test_large_number_of_functions_and_calls():
    """
    Large Scale: Many functions and calls, all get comments.
    """
    project_root, tests_root, test_file = make_file_paths()
    # Generate 100 functions, each with 2 calls
    num_funcs = 100
    test_lines = []
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(num_funcs):
        test_lines.append(f"class TestClass{i}:")
        test_lines.append(f"    def test_func(self):")
        for j in range(2):
            test_lines.append(f"        result{j} = myfunc({i*2+j})")
            test_lines.append(f"        codeflash_output = result{j}")
            # Build invocation id for each call
            inv_id = InvocationId(
                test_module_path="tests.test_file",
                test_class_name=f"TestClass{i}",
                test_function_name="test_func",
                iteration_id=f"{j}_0"
            )
            original_runtimes[inv_id] = [1000 + i*10 + j]
            optimized_runtimes[inv_id] = [500 + i*5 + j]
    test_code = "\n".join(test_lines)
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    test_obj = GeneratedTests(
        generated_original_test_source=test_code,
        behavior_file_path=test_file,
        perf_file_path=test_file
    )
    generated_tests = GeneratedTestsList([test_obj])
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    out_code = result.generated_tests[0].generated_original_test_source

def test_large_number_of_tests_in_list():
    """
    Large Scale: Many test files in GeneratedTestsList, all get processed.
    """
    project_root, tests_root, test_file = make_file_paths()
    num_tests = 50
    generated_tests = []
    original_runtimes = {}
    optimized_runtimes = {}
    for i in range(num_tests):
        code = f"""
class TestClass{i}:
    def test_func(self):
        result = myfunc({i})
        codeflash_output = result
"""
        test_obj = GeneratedTests(
            generated_original_test_source=code,
            behavior_file_path=test_file,
            perf_file_path=test_file
        )
        generated_tests.append(test_obj)
        inv_id = InvocationId(
            test_module_path="tests.test_file",
            test_class_name=f"TestClass{i}",
            test_function_name="test_func",
            iteration_id="0_0"
        )
        original_runtimes[inv_id] = [1000 + i]
        optimized_runtimes[inv_id] = [900 + i]
    test_cfg = TestConfig(tests_root=tests_root, project_root_path=project_root)
    generated_tests_list = GeneratedTestsList(generated_tests)
    codeflash_output = add_runtime_comments_to_generated_tests(
        qualified_name="myfunc",
        test_cfg=test_cfg,
        generated_tests=generated_tests_list,
        original_runtimes=original_runtimes,
        optimized_runtimes=optimized_runtimes
    ); result = codeflash_output
    # All 50 tests should have a comment
    for idx in range(num_tests):
        out_code = result.generated_tests[idx].generated_original_test_source
# codeflash_output is used to check that the output of the original code is the same as that of the optimized code.